import math

import torch
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo
from torch import nn
from torch.nn import Parameter
# import pdb
import numpy as np
import os
# from . import networks
from .base_model import BaseModel
from .cdcn import network
from utils.util import AverageMeter,save_csv,mergetotxt
from utils.util import LD_param as get_res

from PIL import Image
from data.base_dataset import  get_transform
from torchvision import transforms
import cv2

def one_hot(batchsize,numclass,tensor):
    return torch.zeros(batchsize,numclass).to(tensor.device).scatter_(1,tensor.view(1,tensor.shape[0]).long().t(),1)

class CdcnMultiModel(BaseModel):
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.add_argument('--netCDCN', type=str, default='single', help='the architecture situation of model')
        parser.add_argument('--dataIn_1', type=str, default='S0', help='the data type of input1')
        parser.add_argument('--dataIn_2', type=str, default='SD', help='the data type of input2')
        parser.add_argument('--num_classes', type=int, default=2, help='the data type of input2')
        parser.add_argument('--label', type=str, default='label', help='the data type of label')
        if is_train:
            pass
            # parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
            # parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
            # parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')

        return parser
    def display_opt_init(opt):
            # parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        opt.__setattr__('netCDCN','single')
        opt.__setattr__('dataIn_1','S0')
        opt.__setattr__('dataIn_2','SD')
        # opt.__setattr__('cal_mode','linear')
        opt.__setattr__('num_classes',2)
        opt.__setattr__('label','label')
        return opt
    def display_init(self,_dict):
        
        # transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
        # pil_img = Image.fromarray(np.uint8(_dict['img']))
        # _dict[self.opt.inputs] = transform(pil_img)
        # c,h,w = _dict[self.opt.inputs].shape
        # _dict[self.opt.inputs] = _dict[self.opt.inputs].reshape(1,c,h,w)
        _dict['label'] = None
        return _dict
    
    def __init__(self, opt):#basic_conv=networks.Conv2d_cd, theta=0.7 ):   
        BaseModel.__init__(self, opt)
        self.cat_name = None
        self.cat = None
        self.visuale_param()
        # self.set_validation_parameter()
        
        self.model_names=['CDCN']

        self.netCDCN=network.define_CDCN(opt.input_nc, opt.output_nc, opt.ngf, opt.netCDCN, opt.num_classes,
                                         init_type = opt.init_type, init_gain = opt.init_gain,gpu_ids =  self.gpu_ids)
        if self.isTrain:
            # if opt.lambda_identity > 0.0:  # only works when input and output images have the same number of channels
            #     assert(opt.input_nc == opt.output_nc)
            # self.fake_A_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_B_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # define loss functions
            self.last_val_best["Accurate"] = 0
            self.criterion = nn.BCEWithLogitsLoss()

            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer = torch.optim.Adam(self.netCDCN.parameters(), lr=opt.lr, weight_decay=0.00005)
            # self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
            # self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer)
            # self.optimizers.append(self.optimizer_D)
        pass
    def visuale_param(self):
        if self.isTrain:
            self.loss_names=["CDCN_Multi"]
        else:
            self.loss_names += ["Accurate","F_score","APCER","NPCER","ACER"]
            self.TP,self.TN,self.FP,self.FN = {},{},{},{}
            self.TP['all'],self.TN['all'],self.FP['all'],self.FN['all']  = 0,0,0,0
            if self.opt.controller == 'test' or (self.cat_name is not None and not '物理' in self.cat_name.keys()):
                self.TP['物理'],self.TN['物理'],self.FP['物理'],self.FN['物理'], = 0,0,0,0
            self.prediction = []
            self.groud_truth = []
            # self.last_val = {}
            # for key in self.loss_names:
            #     self.last_val[key] = 0
            
    def set_input(self, data):
        # self.input = data[self.opt.in]
        self.dataIn_1 = data[self.opt.dataIn_1]
        self.dataIn_2 = data[self.opt.dataIn_2]
        self.label = data[self.opt.label]
        if self.label is not None:
            self.label_1hot = one_hot(self.label.shape[0],self.opt.num_classes,self.label).to(self.device)
        self.image_paths = None
        
        if 'cat_name' in data.keys() and 'category' in data.keys():
            self.cat_name = data['cat_name']
            self.cat = data['category']
        
        if "isTrain" in data.keys():
            if self.isTrain != data["isTrain"]:
                self.isTrain = data["isTrain"]
                self.visuale_param()
                
        
        
        
        pass
    def forward(self):	    	# x [3, 256, 256]
        self.out = self.netCDCN(self.dataIn_1, self.dataIn_2)
        pass
    def backward(self):
        self.loss_CDCN_Multi = self.criterion(self.out,self.label_1hot)
        self.loss_CDCN_Multi.backward()
    def optimize_parameters(self):
        self.forward()
        self.optimizer.zero_grad()
        self.backward()
        self.optimizer.step()
        pass
    # def cal_score(self):
    def iter_end(self):
        if self.opt.controller == 'test' or not self.isTrain:
            ans = torch.max(self.out,axis=1).indices.cpu()
            # self.loss_pred_acc = (torch.where(ans == self.label)[0].shape[0])  / self.label.shape[0]
            # self.count_mean_pred_acc.update(self.loss_pred_acc,self.label.shape[0])
            # self.loss_mean_pred_acc = self.count_mean_pred_acc.avg
            self.TP['all'] += torch.where((ans==1) &(self.label == 1))[0].shape[0]
            self.TN['all'] += torch.where((ans==0) &(self.label == 0))[0].shape[0]
            self.FP['all'] += torch.where((ans==1) &(self.label == 0))[0].shape[0]
            self.FN['all'] += torch.where((ans==0) &(self.label == 1))[0].shape[0]
            if self.opt.controller == 'test' and self.cat_name is not None and self.cat is not None:
                self.prediction += ans.tolist()
                self.groud_truth += self.label.tolist()
                for key in self.cat_name.keys():
                    if not key in self.TP.keys():
                        self.TP[key]=0
                        self.TN[key]=0
                        self.FP[key]=0
                        self.FN[key]=0
                    self.TP[key] += torch.where((self.cat == self.cat_name[key])&(ans==1) &(self.label == 1))[0].shape[0]
                    self.TN[key] += torch.where((self.cat == self.cat_name[key])&(ans==0) &(self.label == 0))[0].shape[0]
                    self.FP[key] += torch.where((self.cat == self.cat_name[key])&(ans==1) &(self.label == 0))[0].shape[0]
                    self.FN[key] += torch.where((self.cat == self.cat_name[key])&(ans==0) &(self.label == 1))[0].shape[0]
                    if  len(self.cat_name.keys()) > 4 and self.cat >= 0 and self.cat <=4:
                        self.TP['物理'] += torch.where((self.cat == self.cat_name[key])&(ans==1) &(self.label == 1))[0].shape[0]
                        self.TN['物理'] += torch.where((self.cat == self.cat_name[key])&(ans==0) &(self.label == 0))[0].shape[0]
                        self.FP['物理'] += torch.where((self.cat == self.cat_name[key])&(ans==1) &(self.label == 0))[0].shape[0]
                        self.FN['物理'] += torch.where((self.cat == self.cat_name[key])&(ans==0) &(self.label == 1))[0].shape[0]
        
            if not self.isTrain:
                param = get_res(self.TP,self.TN,self.FP,self.FN,2)
            
                self.loss_Accurate = param['Accurate']['all']
                self.loss_F_score = param['F_score']['all']
                self.loss_APCER = param['APCER']['all']
                self.loss_NPCER = param['NPCER']['all']
                self.loss_ACER = param['ACER']['all']
                self.loss_MACER = 1 - self.loss_ACER

    def epoch_end(self):
        # self.count_mean_pred_acc.reset()
        if self.opt.controller =='test':
            param = get_res(self.TP,self.TN,self.FP,self.FN,2)
            csvdir = os.path.join(self.opt.results_dir,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch)
            csvfile = '{}_{}'.format(self.opt.model,'v'+self.opt.version)
            sheet_name = '{}'.format(self.opt.model)
            
            if not '物理' in self.TP.keys():
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all','物理'])
            else:
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all'])

            # roc_plot(self.groud_truth,self.prediction,'test.png','/home/yalin.huang/test/results/'+self.opt.name)
            # mergetotxt(self.groud_truth,self.prediction,'/home/yalin.huang/test/results/'+self.opt.model,'',self.opt.version)
            mergetotxt(self.groud_truth,self.prediction,csvdir,self.opt.model,self.opt.version)
        self.TP,self.TN,self.FP,self.FN = {},{},{},{}
        self.TP['all'],self.TN['all'],self.FP['all'],self.FN['all']  = 0,0,0,0
        if self.opt.controller == 'test' and not '物理' in self.cat_name.keys():
            self.TP['物理'],self.TN['物理'],self.FP['物理'],self.FN['物理'], = 0,0,0,0
        pass
    
    def val_save(self, param):
        if getattr(self,'loss_' + param) == 0 and self.last_val_best[param] == 0:
            if self.last_val_best["Accurate"] < self.loss_Accurate:
                self.last_val_best["Accurate"] = self.loss_Accurate
                return True
        else:
            if self.last_val_best[param] < getattr(self,'loss_' + param):
                self.last_val_best[param] = getattr(self,'loss_' + param)
                return True
        return False
        pass
    
    def display_proc(self,_dict):
        tran = transforms.ToPILImage()
        if len(_dict['boxes'])>0:
            preds = torch.max(self.out,axis=1).indices.cpu()
            for i,pred in enumerate(preds):
                if pred.item():
                    cv2.rectangle(_dict['box_img'], (_dict['boxes'][i][0], _dict['boxes'][i][1]), (_dict['boxes'][i][2], _dict['boxes'][i][3]), (0, 0, 255), 2)
                    text = 'real'
                    cx = _dict['boxes'][i][0]
                    cy = _dict['boxes'][i][1] + 12
                    cv2.putText(_dict['box_img'], text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
                else:
                    cv2.rectangle(_dict['box_img'], (_dict['boxes'][i][0], _dict['boxes'][i][1]), (_dict['boxes'][i][2], _dict['boxes'][i][3]), (255, 0, 0), 2)
                    text = 'fake'
                    cx = _dict['boxes'][i][0]
                    cy = _dict['boxes'][i][1] + 12
                    cv2.putText(_dict['box_img'], text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

            _dict['img'] = _dict['box_img']
        else:
                _dict['img'] = _dict['ori_img']
                
        return _dict